import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from pandas import Timestamp

plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

##### DATA #####
data = {'Task': {0: '大学生创新创业网通知',
                 1: '启航网新闻通知公示',
                 2: '启航行动',
                 3: '双创导航员',
                 4: '双创公开课',
                 5: '科创嘉年华★',
                 6: '办公室秩序维持',
                 7: '年鉴制作',
                 8: '赛事认证★',
                 9: '学院负责人培训',
                 10: '竞赛通知发布',
                 11: '“挑战杯”校赛★',
                 12: '“五四杯”（创新赛道）★',
                 13: '“启航杯”★'},

        'Department': {0: '媒体部',
                       1: '媒体部',
                       2: '科普培训部',
                       3: '科普培训部',
                       4: '科普培训部',
                       5: '科普培训部',
                       6: '办公室',
                       7: '竞赛监察部',
                       8: '竞赛监察部',
                       9: '竞赛监察部',
                       10: '竞赛监察部',
                       11: '校外竞赛部',
                       12: '校内竞赛部',
                       13: '校内竞赛部'},

        'Start': {0: Timestamp('2025-06-01 00:00:00'),
                  1: Timestamp('2025-06-01 00:00:00'),
                  2: Timestamp('2025-06-01 00:00:00'),
                  3: Timestamp('2025-06-01 00:00:00'),
                  4: Timestamp('2025-06-01 00:00:00'),
                  5: Timestamp('2025-09-15 00:00:00'),
                  6: Timestamp('2025-06-01 00:00:00'),
                  7: Timestamp('2025-12-15 00:00:00'),
                  8: Timestamp('2025-06-15 00:00:00'),
                  9: Timestamp('2025-07-15 00:00:00'),
                  10: Timestamp('2025-06-01 00:00:00'),
                  11: Timestamp('2026-02-15 00:00:00'),
                  12: Timestamp('2025-11-15 00:00:00'),
                  13: Timestamp('2025-10-15 00:00:00')},

        'End': {0: Timestamp('2026-06-01 00:00:00'),
                1: Timestamp('2026-06-01 00:00:00'),
                2: Timestamp('2026-06-01 00:00:00'),
                3: Timestamp('2026-06-01 00:00:00'),
                4: Timestamp('2026-06-01 00:00:00'),
                5: Timestamp('2025-10-15 00:00:00'),
                6: Timestamp('2026-06-01 00:00:00'),
                7: Timestamp('2026-02-15 00:00:00'),
                8: Timestamp('2025-09-15 00:00:00'),
                9: Timestamp('2025-08-15 00:00:00'),
                10: Timestamp('2026-06-01 00:00:00'),
                11: Timestamp('2026-05-15 00:00:00'),
                12: Timestamp('2025-12-15 00:00:00'),
                13: Timestamp('2025-11-15 00:00:00')},

        'Completion': {0: 0.0,
                       1: 0.0,
                       2: 0.0,
                       3: 0.0,
                       4: 0.0,
                       5: 0.0,
                       6: 0.08,
                       7: 0.0,
                       8: 0.3,
                       9: 0.0,
                       10: 0.08,
                       11: 0.0,
                       12: 0.0,
                       13: 0.0}}

##### DATA PREP #####
df = pd.DataFrame(data)

# project start date
proj_start = df.Start.min()

# number of days from project start to task start
df['start_num'] = (df.Start - proj_start).dt.days

# number of days from project start to end of tasks
df['end_num'] = (df.End - proj_start).dt.days

# days between start and end of each task
df['days_start_to_end'] = df.end_num - df.start_num

# days between start and current progression of each task
df['current_num'] = (df.days_start_to_end * df.Completion)


# create a column with the color for each department
def color(row):
    c_dict = {'校内竞赛部': '#E64646', '校外竞赛部': '#E69646', '竞赛监察部': '#34D05C', '办公室': "#F4E869", '科普培训部': '#34D0C3', '媒体部': '#3475D0'}
    return c_dict[row['Department']]


df['color'] = df.apply(color, axis=1)

##### PLOT #####
fig, (ax, ax1) = plt.subplots(2, figsize=(16, 6), gridspec_kw={'height_ratios': [6, 1]}, facecolor="#FFFFFF")
ax.set_facecolor("#FFFFFF")
ax1.set_facecolor("#FFFFFF")
# bars

ax.barh(df.Task, df.current_num, left=df.start_num, color=df.color)

for idx, row in df.iterrows():
    # 判断任务是否为重点
    if '★' in row.Task:
        color_edge = 'red'
        line_width = 3
    else:
        color_edge = 'black'
        line_width = 0

    ax.barh(row.Task, row.days_start_to_end, left=row.start_num, color=row.color, alpha=0.5, edgecolor=color_edge, linewidth=line_width)


# ax.barh(df.Task, df.days_start_to_end, left=df.start_num, color=df.color, alpha=0.5)

for idx, row in df.iterrows():
    ax.text(row.end_num + 0.1, idx, f"{int(row.Completion * 100)}%", va='center', alpha=0.8, color='black')
    ax.text(row.start_num - 0.1, idx, row.Task, va='center', ha='right', alpha=0.8, color='black')

# grid lines
ax.set_axisbelow(True)
ax.xaxis.grid(color='k', linestyle='dashed', alpha=0.4, which='both')

# ticks
xticks = np.arange(0, df.end_num.max() + 1, 30)
xticks_labels = pd.date_range(proj_start, end=df.End.max()).strftime("%m/%d")
xticks_minor = np.arange(0, df.end_num.max() + 1, 10)
ax.set_xticks(xticks)
ax.set_xticks(xticks_minor, minor=True)
ax.set_xticklabels(xticks_labels[::30], color='black')
ax.set_yticks([])

plt.setp([ax.get_xticklines()], color='black')

# align x axis
ax.set_xlim(0, df.end_num.max())

# remove spines
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['left'].set_position(('outward', 10))
ax.spines['top'].set_visible(False)
ax.spines['bottom'].set_color('black')

plt.suptitle('校学生科技协会2025-2026工作规划', color='black')

##### LEGENDS #####
legend_elements = [Patch(facecolor='#E64646', label='校内竞赛部'),
                   Patch(facecolor='#E69646', label='校外竞赛部'),
                   Patch(facecolor='#34D05C', label='竞赛监察部'),
                   Patch(facecolor='#F4E869', label='办公室'),
                   Patch(facecolor='#34D0C3', label='科普培训部'),
                   Patch(facecolor='#3475D0', label='媒体部')]

legend = ax1.legend(handles=legend_elements, loc='upper center', ncol=5, frameon=False)
plt.setp(legend.get_texts(), color='black')

# clean second axis
ax1.spines['right'].set_visible(False)
ax1.spines['left'].set_visible(False)
ax1.spines['top'].set_visible(False)
ax1.spines['bottom'].set_visible(False)
ax1.set_xticks([])
ax1.set_yticks([])

plt.savefig('gantt.png', facecolor="#FFFFFF")